#!/usr/bin/env python3

import xarray as xr
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
from cartopy.util import add_cyclic

f = xr.open_dataset('adv_sr.720x360.h0.nc')

fig = plt.figure(figsize=(10, 5))

q, lon, lat = add_cyclic(f.q1[3,0,:,:].values, f.lon, f.lat)
ax1 = fig.add_subplot(1, 2, 1, projection=ccrs.Stereographic(central_latitude=90))
plt.colorbar(ax1.contourf(lon, lat, q, transform=ccrs.PlateCarree(), cmap='jet'), orientation='horizontal')
ax1.contour(lon, lat, q, transform=ccrs.PlateCarree(), colors='k')
ax1.set_extent([-180, 180, 70, 90], ccrs.PlateCarree())
ax1.set_title(f'Day 6\nMin: {f.q1[3,0,:,:].min().values:.4f}, Max: {f.q1[3,0,:,:].max().values:.2f}')

q, lon, lat = add_cyclic(f.q1[12,0,:,:].values, f.lon, f.lat)
ax2 = fig.add_subplot(1, 2, 2, projection=ccrs.PlateCarree(central_longitude=180))
plt.colorbar(ax2.contourf(lon, lat, q, transform=ccrs.PlateCarree(), cmap='jet'), orientation='horizontal')
ax2.contour(lon, lat, q, transform=ccrs.PlateCarree(), colors='k')
ax2.set_extent([250, 290, -20, 20], ccrs.PlateCarree())
ax2.set_title(f'Day 12\nMin: {f.q1[12,0,:,:].min().values:.4f}, Max: {f.q1[12,0,:,:].max().values:.2f}')

plt.tight_layout()
plt.show()
